from django.core.management.base import BaseCommand
from accounts.models import Institution, Laboratory


class Command(BaseCommand):
    help = '初始化机构和实验室数据'

    def handle(self, *args, **options):
        # 创建机构
        institutions_data = [
            {'code': 'CAU', 'name': '中国农业大学'},
            {'code': 'OTHER', 'name': '其他机构'},
        ]

        for inst_data in institutions_data:
            institution, created = Institution.objects.get_or_create(
                code=inst_data['code'],
                defaults={'name': inst_data['name']}
            )
            if created:
                self.stdout.write(
                    self.style.SUCCESS(f'成功创建机构: {institution.name}')
                )
            else:
                self.stdout.write(
                    self.style.WARNING(f'机构已存在: {institution.name}')
                )

        # 创建实验室
        laboratories_data = [
            {
                'institution_code': 'CAU',
                'code': 'SASL',
                'name': '智慧农业系统实验室'
            },
            {
                'institution_code': 'CAU',
                'code': 'PDRL',
                'name': '植物疾病识别实验室'
            },
            {
                'institution_code': 'OTHER',
                'code': 'OTHER',
                'name': '其他实验室'
            },
        ]

        for lab_data in laboratories_data:
            try:
                institution = Institution.objects.get(code=lab_data['institution_code'])
                laboratory, created = Laboratory.objects.get_or_create(
                    institution=institution,
                    code=lab_data['code'],
                    defaults={'name': lab_data['name']}
                )
                if created:
                    self.stdout.write(
                        self.style.SUCCESS(f'成功创建实验室: {laboratory.name}')
                    )
                else:
                    self.stdout.write(
                        self.style.WARNING(f'实验室已存在: {laboratory.name}')
                    )
            except Institution.DoesNotExist:
                self.stdout.write(
                    self.style.ERROR(f'机构不存在: {lab_data["institution_code"]}')
                )

        self.stdout.write(
            self.style.SUCCESS('机构和实验室初始化完成！')
        )